{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Semantic Search using the Inference API with the AlibabaCloudSearch service\n",
    "Learn how to use the [Inference API](https://www.elastic.co/guide/en/elasticsearch/reference/current/inference-apis.html) for semantic search."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Requirements\n",
    "\n",
    "For this example, you will need:\n",
    "\n",
    "- An Elastic deployment:\n",
    "  - we'll be using [Alibaba Elasticsearch](https://www.aliyun.com/product/bigdata/elasticsearch) for this example.\n",
    "- Elasticsearch 8.16 or above\n",
    "- A valid API key for the [AlibabaCloud AI Search](https://help.aliyun.com/zh/open-search/search-platform/user-guide/activate-services-and-create-api-key?spm=a2c4g.11186623.0.0.141a2105fiqCEw) is required to use the Inference API with\n",
    "the AlibabaCloud AI Search service."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Install packages and connect with Elasticsearch Client\n",
    "\n",
    "To get started, we'll need to connect to our Elastic deployment using the Python client (version 8.15.0 or above).\n",
    "\n",
    "First we need to `pip` install the following packages:\n",
    "\n",
    "- `elasticsearch`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install elasticsearch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we need to import the modules we need. \n",
    "\n",
    "🔐 NOTE: getpass enables us to securely prompt the user for credentials without echoing them to the terminal, or storing it in memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from elasticsearch import Elasticsearch, helpers\n",
    "from urllib.request import urlopen\n",
    "from getpass import getpass\n",
    "import json\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can instantiate the Python Elasticsearch client.\n",
    "\n",
    "First we prompt the user for their user name, password and elastic host.\n",
    "Then we create a `client` object that instantiates an instance of the `Elasticsearch` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "ELASTIC_USER = getpass(\"ELASTIC USER: \")\n",
    "\n",
    "ELASTIC_PASSWORD = getpass(\"ELASTIC PASSWORD: \")\n",
    "\n",
    "# e.g.\n",
    "# hosts = \"http://es-xxx.elasticsearch.aliyuncs.com:9200\"\n",
    "hosts = getpass(\"Host: \")\n",
    "\n",
    "# Create the client instance\n",
    "client = Elasticsearch(\n",
    "    # For local development\n",
    "    # hosts=[\"http://localhost:9200\"]\n",
    "    hosts,\n",
    "    basic_auth=(ELASTIC_USER, ELASTIC_PASSWORD),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Confirm that the client has connected with this test:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(client.info())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Refer to [the documentation](https://www.elastic.co/guide/en/elasticsearch/client/python-api/current/connecting.html#connect-self-managed-new) to learn how to connect to a self-managed deployment.\n",
    "\n",
    "Read [this page](https://www.elastic.co/guide/en/elasticsearch/client/python-api/current/connecting.html#connect-self-managed-new) to learn how to connect using API keys."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"create-the-inference-endpoint\"></a>\n",
    "## Create the inference endpoint\n",
    "\n",
    "Let's create the inference endpoint by using the [Create inference API](https://www.elastic.co/guide/en/elasticsearch/reference/current/put-inference-api.html).\n",
    "\n",
    "You'll need a valid API key for AlibabaCloud AI Search for this. You can create one by referring to the [API keys management](https://help.aliyun.com/zh/open-search/search-platform/user-guide/api-keys-management)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "API_KEY = getpass(\"Enter AlibabaCloud AI Search API key:  \")\n",
    "\n",
    "# Please enter your AlibabaCloud AI Search host here\n",
    "# e.g.\n",
    "# HOST = \"default-xxx.opensearch.aliyuncs.com\"\n",
    "HOST = \"\"\n",
    "\n",
    "# Please enter your AlibabaCloud AI Search workspace name here\n",
    "# e.g.\n",
    "# WORKSPACE = \"default\"\n",
    "WORKSPACE = \"\"\n",
    "\n",
    "client.inference.put(\n",
    "    task_type=\"text_embedding\",\n",
    "    inference_id=\"os-embeddings-notebook-test\",\n",
    "    body={\n",
    "        \"service\": \"alibabacloud-ai-search\",\n",
    "        \"service_settings\": {\n",
    "            \"api_key\": API_KEY,\n",
    "            \"service_id\": \"ops-text-embedding-001\",\n",
    "            \"host\": HOST,\n",
    "            \"workspace\": WORKSPACE,\n",
    "        },\n",
    "    },\n",
    ")\n",
    "\n",
    "client.inference.put(\n",
    "    task_type=\"sparse_embedding\",\n",
    "    inference_id=\"os-sparse-embeddings-notebook-test\",\n",
    "    body={\n",
    "        \"service\": \"alibabacloud-ai-search\",\n",
    "        \"service_settings\": {\n",
    "            \"api_key\": API_KEY,\n",
    "            \"service_id\": \"ops-text-sparse-embedding-001\",\n",
    "            \"host\": HOST,\n",
    "            \"workspace\": WORKSPACE,\n",
    "        },\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an ingest pipeline with an inference processor\n",
    "\n",
    "Create an ingest pipeline with an inference processor by using the [`put_pipeline`](https://www.elastic.co/guide/en/elasticsearch/reference/master/put-pipeline-api.html) method. Reference the inference endpoint created above as the `model_id` to infer against the data that is being ingested in the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ObjectApiResponse({'acknowledged': True})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.ingest.put_pipeline(\n",
    "    id=\"alibaba-embeddings-notebook-test\",\n",
    "    description=\"Ingest pipeline for alibaba inference notebook test.\",\n",
    "    processors=[\n",
    "        {\n",
    "            \"inference\": {\n",
    "                \"model_id\": \"os-embeddings-notebook-test\",\n",
    "                \"input_output\": {\n",
    "                    \"input_field\": \"plot\",\n",
    "                    \"output_field\": \"plot_embedding\",\n",
    "                },\n",
    "            }\n",
    "        },\n",
    "        {\n",
    "            \"inference\": {\n",
    "                \"model_id\": \"os-sparse-embeddings-notebook-test\",\n",
    "                \"input_output\": {\n",
    "                    \"input_field\": \"plot\",\n",
    "                    \"output_field\": \"plot_sparse_embedding\",\n",
    "                },\n",
    "            }\n",
    "        },\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's note a few important parameters from that API call:\n",
    "\n",
    "- `inference`: A processor that performs inference using a machine learning model.\n",
    "- `model_id`: Specifies the ID of the inference endpoint to be used. In this example, the model ID is set to `cohere_embeddings`.\n",
    "- `input_output`: Specifies input and output fields.\n",
    "- `input_field`: Field name from which the `dense_vector` and `sparse_vector`representation is created.\n",
    "- `output_field`:  Field name which contains inference results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create index\n",
    "\n",
    "The mapping of the destination index – the index that contains the embeddings and sparse_embeddings that the model will create based on your input text – must be created. The destination index must have a field with the [dense_vector](https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html) field type and a field with the [sparse_vector](https://www.elastic.co/guide/en/elasticsearch/reference/current/sparse-vector.html) field type to index the output of the AlibabaCloud AI Search model.\n",
    "\n",
    "Let's create an index named `alibaba-text-embeddings-notebook-test` with the mappings we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ObjectApiResponse({'acknowledged': True, 'shards_acknowledged': True, 'index': 'alibaba-text-embeddings-notebook-test'})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.indices.delete(\n",
    "    index=\"alibaba-text-embeddings-notebook-test\", ignore_unavailable=True\n",
    ")\n",
    "client.indices.create(\n",
    "    index=\"alibaba-text-embeddings-notebook-test\",\n",
    "    settings={\"index\": {\"default_pipeline\": \"alibaba-embeddings-notebook-test\"}},\n",
    "    mappings={\n",
    "        \"properties\": {\n",
    "            \"plot\": {\"type\": \"text\"},\n",
    "            \"plot_embedding\": {\"type\": \"dense_vector\", \"dims\": 1536, \"index\": \"true\"},\n",
    "            \"plot_sparse_embedding\": {\"type\": \"sparse_vector\"},\n",
    "        }\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Insert Documents\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done indexing documents into `alibaba-text-embeddings-notebook-test` index!\n"
     ]
    }
   ],
   "source": [
    "url = \"https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/notebooks/search/movies.json\"\n",
    "response = urlopen(url)\n",
    "\n",
    "# Load the response data into a JSON object\n",
    "data_json = json.loads(response.read())\n",
    "\n",
    "# Prepare the documents to be indexed\n",
    "documents = []\n",
    "for doc in data_json:\n",
    "    documents.append(\n",
    "        {\n",
    "            \"_index\": \"alibaba-text-embeddings-notebook-test\",\n",
    "            \"_source\": doc,\n",
    "        }\n",
    "    )\n",
    "\n",
    "# Use helpers.bulk to index\n",
    "helpers.bulk(client, documents)\n",
    "\n",
    "print(\"Done indexing documents into `alibaba-text-embeddings-notebook-test` index!\")\n",
    "time.sleep(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Semantic search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### knn search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Score: 0.6912112\n",
      "Title: Pulp Fiction\n",
      "Plot: The lives of two mob hitmen, a boxer, a gangster and his wife, and a pair of diner bandits intertwine in four tales of violence and redemption.\n",
      "\n",
      "Score: 0.66019344\n",
      "Title: The Usual Suspects\n",
      "Plot: A sole survivor tells of the twisty events leading up to a horrific gun battle on a boat, which began when five criminals met at a seemingly random police lineup.\n",
      "\n",
      "Score: 0.64070797\n",
      "Title: The Dark Knight\n",
      "Plot: When the menace known as the Joker wreaks havoc and chaos on the people of Gotham, Batman must accept one of the greatest psychological and physical tests of his ability to fight injustice.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "response = client.search(\n",
    "    index=\"alibaba-text-embeddings-notebook-test\",\n",
    "    size=3,\n",
    "    knn={\n",
    "        \"field\": \"plot_embedding\",\n",
    "        \"query_vector_builder\": {\n",
    "            \"text_embedding\": {\n",
    "                \"model_id\": \"os-embeddings-notebook-test\",\n",
    "                \"model_text\": \"Fighting movie\",\n",
    "            }\n",
    "        },\n",
    "        \"k\": 10,\n",
    "        \"num_candidates\": 100,\n",
    "    },\n",
    ")\n",
    "\n",
    "for hit in response[\"hits\"][\"hits\"]:\n",
    "    doc_id = hit[\"_id\"]\n",
    "    score = hit[\"_score\"]\n",
    "    title = hit[\"_source\"][\"title\"]\n",
    "    plot = hit[\"_source\"][\"plot\"]\n",
    "    print(f\"Score: {score}\\nTitle: {title}\\nPlot: {plot}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### sparse vector search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Score: 0.1208359\n",
      "Title: The Godfather\n",
      "Plot: An organized crime dynasty's aging patriarch transfers control of his clandestine empire to his reluctant son.\n",
      "\n",
      "Score: 0.050445557\n",
      "Title: Goodfellas\n",
      "Plot: The story of Henry Hill and his life in the mob, covering his relationship with his wife Karen Hill and his mob partners Jimmy Conway and Tommy DeVito in the Italian-American crime syndicate.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "response = client.search(\n",
    "    index=\"alibaba-text-embeddings-notebook-test\",\n",
    "    query={\n",
    "        \"sparse_vector\": {\n",
    "            \"field\": \"plot_sparse_embedding\",\n",
    "            \"inference_id\": \"os-sparse-embeddings-notebook-test\",\n",
    "            \"query\": \"crime dynasty\",\n",
    "        }\n",
    "    },\n",
    ")\n",
    "\n",
    "for hit in response[\"hits\"][\"hits\"]:\n",
    "    doc_id = hit[\"_id\"]\n",
    "    score = hit[\"_score\"]\n",
    "    title = hit[\"_source\"][\"title\"]\n",
    "    plot = hit[\"_source\"][\"plot\"]\n",
    "    print(f\"Score: {score}\\nTitle: {title}\\nPlot: {plot}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE:** The value of `model_id` in the `query_vector_builder` and the value of `inference_id` in the `sparse_vector` must match the value of `inference_id` you created in the [first step](#create-the-inference-endpoint)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
